/*
 * To change this template, choose Tools | Templates
 * and open the template in the editor.
 */
package pt.utl.ist.isr.livetrack.analysis;

import static java.lang.Math.*;
import java.awt.image.BufferedImage;
import java.awt.image.ColorModel;
import java.awt.image.Raster;
import java.awt.image.WritableRaster;
import java.util.ArrayList;
import java.util.Arrays;
import org.apache.log4j.Logger;
import pt.utl.ist.isr.livetrack.server.interfaces.Blob;
import pt.utl.ist.isr.livetrack.server.impl.BlobImpl;
import pt.utl.ist.isr.livetrack.analysis.util.UnionFind;

/**
 *
 * @author David Miguel Antunes <davidmiguel [ at ] antunes.net>
 */
public class ImageUtils {

    private static final Logger logger = Logger.getLogger(ImageUtils.class.getName());
    private final int width;
    private final int height;
    private final int nComp;
    private int[] tmpWidthHeightIntArray;
    private UnionFind unionFind;

    public ImageUtils(int width, int height, int nComponents) {
        this.width = width;
        this.height = height;
        this.nComp = nComponents;
    }

    private int[] getTmpWidthHeightIntArray() {
        if (tmpWidthHeightIntArray == null) {
            tmpWidthHeightIntArray = new int[width * height];
        }
        return tmpWidthHeightIntArray;
    }

    private UnionFind getUnionFind() {
        if (unionFind == null) {
            unionFind = new UnionFind((int) (width * height * 1));
        }
        return unionFind;
    }

    public boolean valid(int width, int height, int nComponents) {
        return this.width == width && this.height == height && this.nComp == nComponents;
    }

    // ==== Start Utility Methods ====
    public int[] newImage() {
        return new int[width * height * nComp];
    }

    public int[] newIntensityImage() {
        return new int[width * height * 1];
    }

    public boolean[] newBinaryImage() {
        return new boolean[width * height];
    }

    public BufferedImage arrayToBufferedImage(int[] img) {
        BufferedImage i = new BufferedImage(width, height, nComp == 3 ? BufferedImage.TYPE_3BYTE_BGR : BufferedImage.TYPE_USHORT_GRAY);

        WritableRaster r = i.getRaster();
        if (nComp == 3) {
            for (int l = 0; l < height; l++) {
                for (int c = 0; c < width; c++) {
                    r.setPixel(c, l, new int[]{
                                img[l * width * nComp + c * nComp + 0],
                                img[l * width * nComp + c * nComp + 1],
                                img[l * width * nComp + c * nComp + 2]});
                }
            }
        } else if (nComp == 1) {
            for (int l = 0; l < height; l++) {
                for (int c = 0; c < width; c++) {
                    r.setPixel(c, l, new int[]{img[l * width + c] * 256});
                }
            }
        } else {
            throw new RuntimeException("Cannot convert array to image");
        }

        return i;
    }

    public BufferedImage arrayToBufferedImage(boolean[] img) {
        BufferedImage i = new BufferedImage(width, height, BufferedImage.TYPE_USHORT_GRAY);

        WritableRaster r = i.getRaster();
        for (int l = 0; l < height; l++) {
            for (int c = 0; c < width; c++) {
                r.setPixel(c, l, new int[]{img[l * width + c] ? 65535 : 0});
            }
        }

        return i;
    }

    public int[] bufferedImageToArray(BufferedImage img) {

        Raster r = img.getRaster();
        int[] i = newImage();

        final int[] pixel = new int[3];
        for (int l = 0; l < height; l++) {
            for (int c = 0; c < width; c++) {
                r.getPixel(c, l, pixel);
                for (int j = 0; j < nComp; j++) {
                    i[l * width * nComp + c * nComp + j] = pixel[j];
                }
            }
        }

        return i;
    }

    public void bufferedImageToArray(BufferedImage img, int[] array) {

        Raster r = img.getRaster();

        final int[] pixel = new int[3];
        for (int l = 0; l < height; l++) {
            for (int c = 0; c < width; c++) {
                r.getPixel(c, l, pixel);
                for (int j = 0; j < nComp; j++) {
                    array[l * width * nComp + c * nComp + j] = pixel[j];
                }
            }
        }
    }

    public static BufferedImage deepCopy(BufferedImage bi) {
        ColorModel cm = bi.getColorModel();
        boolean isAlphaPremultiplied = cm.isAlphaPremultiplied();
        WritableRaster raster = bi.copyData(null);
        return new BufferedImage(cm, raster, isAlphaPremultiplied, null);
    }

    public int[][][] deepCopy(int[][][] orig) {
        int[][][] dest = new int[orig.length][orig[0].length][orig[0][0].length];
        for (int i = 0; i < dest.length; i++) {
            for (int j = 0; j < dest[0].length; j++) {
                System.arraycopy(orig[i][j], 0, dest[i][j], 0, dest[0][0].length);
            }
        }
        return dest;
    }

    public int[] deepCopy(int[] orig) {
        int[] dest = new int[orig.length];
        System.arraycopy(orig, 0, dest, 0, dest.length);
        return dest;
    }

    public boolean[] deepCopy(boolean[] orig) {
        boolean[] dest = new boolean[orig.length];
        System.arraycopy(orig, 0, dest, 0, dest.length);
        return dest;
    }

    private void print(boolean[] img) {
        for (int i = 0; i < height; i++) {
            for (int j = 0; j < width; j++) {
                System.out.print(img[i * width + j] ? '1' : '0');
            }
            System.out.println("");
        }
    }
    // ==== End Utility Methods ====

    // ==== Start Algorithmic Methods ====
    public void dilate(
            final boolean[] in,
            final boolean[] out,
            final boolean[][] kernel,
            final boolean distanceOpt) {
        if (kernel.length % 2 != 1 || kernel[0].length % 2 != 1) {
            throw new RuntimeException("Kernel must have an uneven size");
        }

        final int verticalMargin = kernel.length / 2;
        final int horizontalMargin = kernel[0].length / 2;

        System.arraycopy(in, 0, out, 0, out.length);

        int[] distance = null;
        int radius = max(kernel.length, kernel[0].length) / 2;
        if (distanceOpt) {
            distance = distanceMatrix(in, true, radius + 2);
        }

        for (int l = 0; l < height; l++) {
            for (int c = 0; c < width; c++) {
                if (!in[l * width + c] && (!distanceOpt || distance[l * width + c] <= radius)) {
                    // Background:
                    kernelSearchLoop:
                    for (int l2 = Math.max(0, l - verticalMargin); l2 <= Math.min(height - 1, l + verticalMargin); l2++) {
                        for (int c2 = Math.max(0, c - horizontalMargin); c2 <= Math.min(width - 1, c + horizontalMargin); c2++) {
                            if (kernel[l2 + verticalMargin - l][c2 + horizontalMargin - c] && in[l2 * width + c2]) {
                                out[l * width + c] = true;
                                break kernelSearchLoop;
                            }
                        }
                    }
                }
            }
        }
    }

    public void erode(
            final boolean[] in,
            final boolean[] out,
            final boolean[][] kernel,
            final boolean distanceOpt) {
        if (kernel.length % 2 != 1 || kernel[0].length % 2 != 1) {
            throw new RuntimeException("Kernel must have an uneven size");
        }

        final int verticalMargin = kernel.length / 2;
        final int horizontalMargin = kernel[0].length / 2;

        System.arraycopy(in, 0, out, 0, out.length);

        int[] distance = null;
        int radius = max(kernel.length, kernel[0].length) / 2;
        if (distanceOpt) {
            distance = distanceMatrix(in, false, radius + 2);
        }

        for (int l = 0; l < height; l++) {
            for (int c = 0; c < width; c++) {
                if (in[l * width + c] && (!distanceOpt || distance[l * width + c] <= radius)) {
                    // Foreground:
                    kernelSearchLoop:
                    for (int l2 = Math.max(0, l - verticalMargin); l2 <= Math.min(height - 1, l + verticalMargin); l2++) {
                        for (int c2 = Math.max(0, c - horizontalMargin); c2 <= Math.min(width - 1, c + horizontalMargin); c2++) {
                            if (kernel[l2 + verticalMargin - l][c2 + horizontalMargin - c] && !in[l2 * width + c2]) {
                                out[l * width + c] = false;
                                break kernelSearchLoop;
                            }
                        }
                    }
                }
            }
        }
    }

    /**
     * @param marked Will be filled with the component number (zero-based) or -1
     * if the pixel does not belong to any component.
     */
    public int extractConnectedComponents(
            final boolean[] bin,
            // output
            final int[] marked) {
        int lblgen = 1;

        Arrays.fill(marked, 0);

        UnionFind uf = getUnionFind();
        uf.reset();

        int[] neighbors = new int[4];
        for (int l = 0; l < height; l++) {
            for (int c = 0; c < width; c++) {
                if (bin[l * width + c]) {

                    neighbors[0] = c == 0 ? 0 : marked[l * width + (c - 1)];
                    neighbors[1] = l == 0 || c == 0 ? 0 : marked[(l - 1) * width + (c - 1)];
                    neighbors[2] = l == 0 ? 0 : marked[(l - 1) * width + c];
                    neighbors[3] = l == 0 | c == (width - 1) ? 0 : marked[(l - 1) * width + (c + 1)];

                    if (neighbors[0] == neighbors[1]
                            && neighbors[1] == neighbors[2]
                            && neighbors[2] == neighbors[3]) {
                        if (neighbors[0] == 0) {
                            marked[l * width + c] = lblgen;
                            lblgen++;
                        } else {
                            marked[l * width + c] = neighbors[0];
                        }
                    } else {
                        int i = 0;
                        for (; i < neighbors.length; i++) {
                            if (neighbors[i] != 0) {
                                marked[l * width + c] = neighbors[i];
                                break;
                            }
                        }
                        for (int j = i + 1; j < neighbors.length; j++) {
                            if (neighbors[j] != 0) {
                                uf.unite(neighbors[i], neighbors[j]);
                            }
                        }
                    }
                }
            }
        }

        boolean[] usedBlobIds = new boolean[lblgen];
        int[] groupIdToBlobId = new int[lblgen];
        int[] markToFinalId = new int[lblgen];

        int nComponents = 0;
        for (int i = 0; i < width * height; i++) {
            int comp = marked[i];
            if (comp == 0) {
                markToFinalId[comp] = -1;
            } else {
                int group = uf.find(marked[i]);
                if (!usedBlobIds[group]) {
                    usedBlobIds[group] = true;
                    groupIdToBlobId[group] = nComponents++;
                }
                markToFinalId[comp] = groupIdToBlobId[group];
            }

        }

        for (int l = 0; l < height; l++) {
            for (int c = 0; c < width; c++) {
                marked[l * width + c] = markToFinalId[marked[l * width + c]];
            }
        }

        return nComponents;
    }

    public ArrayList<Blob> connectedComponentsToBlobs(
            final int[] marked,
            final int nLabels,
            final String camId) {
        ArrayList<Blob> blobs = new ArrayList<Blob>();
        int[] corners = new int[4 * nLabels];
        int[] sizes = new int[nLabels];

        for (int i = 0; i < nLabels; i++) {
            corners[i * 4 + 0] = width;
            corners[i * 4 + 1] = height;
        }

        for (int l = 0; l < height; l++) {
            for (int c = 0; c < width; c++) {
                int lbl = marked[l * width + c];
                if (lbl != -1) {
                    if (c < corners[4 * lbl + 0]) {
                        corners[4 * lbl + 0] = c;
                    }
                    if (l < corners[4 * lbl + 1]) {
                        corners[4 * lbl + 1] = l;
                    }
                    if (c > corners[4 * lbl + 2]) {
                        corners[4 * lbl + 2] = c;
                    }
                    if (l > corners[4 * lbl + 3]) {
                        corners[4 * lbl + 3] = l;
                    }
                    sizes[lbl]++;
                }
            }
        }

        for (int i = 0; i < nLabels; i++) {
            BlobImpl blob = new BlobImpl(
                    corners[4 * i + 0],
                    corners[4 * i + 1],
                    corners[4 * i + 2],
                    corners[4 * i + 3],
                    sizes[i],
                    camId);
            blobs.add(blob);
        }
        return blobs;
    }

    private void fill(
            final int[] distance,
            final int maxDist,
            final int ll,
            final int cc,
            final int direction) {
        int d = distance[ll * width + cc];
        int l, c;

        switch (direction) {
            // UP:
            case 0:
                if (ll > 0) {
                    l = -1;
                    c = -1;
                    if (cc > 0 && distance[(ll + l) * width + (cc + c)] > d + 1 && d + 1 <= maxDist) {
                        distance[(ll + l) * width + (cc + c)] = d + 1;
                        fill(distance, maxDist, ll + l, cc + c, direction);
                    }
                    c = 0;
                    if (distance[(ll + l) * width + (cc + c)] > d + 1 && d + 1 <= maxDist) {
                        distance[(ll + l) * width + (cc + c)] = d + 1;
                        fill(distance, maxDist, ll + l, cc + c, direction);
                    }
                    c = 1;
                    if (cc < width - 1 && distance[(ll + l) * width + (cc + c)] > d + 1 && d + 1 <= maxDist) {
                        distance[(ll + l) * width + (cc + c)] = d + 1;
                        fill(distance, maxDist, ll + l, cc + c, direction);
                    }
                }
                break;
            // BOTTOM:
            case 2:
                if (ll < height - 1) {
                    l = 1;
                    c = -1;
                    if (cc > 0 && distance[(ll + l) * width + (cc + c)] > d + 1 && d + 1 <= maxDist) {
                        distance[(ll + l) * width + (cc + c)] = d + 1;
                        fill(distance, maxDist, ll + l, cc + c, direction);
                    }
                    c = 0;
                    if (distance[(ll + l) * width + (cc + c)] > d + 1 && d + 1 <= maxDist) {
                        distance[(ll + l) * width + (cc + c)] = d + 1;
                        fill(distance, maxDist, ll + l, cc + c, direction);
                    }
                    c = 1;
                    if (cc < width - 1 && distance[(ll + l) * width + (cc + c)] > d + 1 && d + 1 <= maxDist) {
                        distance[(ll + l) * width + (cc + c)] = d + 1;
                        fill(distance, maxDist, ll + l, cc + c, direction);
                    }
                }
                break;
            // RIGHT:
            case 1:
                if (cc < width - 1) {
                    c = 1;
                    l = -1;
                    if (ll > 0 && distance[(ll + l) * width + (cc + c)] > d + 1 && d + 1 <= maxDist) {
                        distance[(ll + l) * width + (cc + c)] = d + 1;
                        fill(distance, maxDist, ll + l, cc + c, direction);
                    }
                    l = 0;
                    if (distance[(ll + l) * width + (cc + c)] > d + 1 && d + 1 <= maxDist) {
                        distance[(ll + l) * width + (cc + c)] = d + 1;
                        fill(distance, maxDist, ll + l, cc + c, direction);
                    }
                    l = 1;
                    if (ll < height - 1 && distance[(ll + l) * width + (cc + c)] > d + 1 && d + 1 <= maxDist) {
                        distance[(ll + l) * width + (cc + c)] = d + 1;
                        fill(distance, maxDist, ll + l, cc + c, direction);
                    }
                }
                break;
            // LEFT:
            case 3:
                if (cc > 0) {
                    c = -1;
                    l = -1;
                    if (ll > 0 && distance[(ll + l) * width + (cc + c)] > d + 1 && d + 1 <= maxDist) {
                        distance[(ll + l) * width + (cc + c)] = d + 1;
                        fill(distance, maxDist, ll + l, cc + c, direction);
                    }
                    l = 0;
                    if (distance[(ll + l) * width + (cc + c)] > d + 1 && d + 1 <= maxDist) {
                        distance[(ll + l) * width + (cc + c)] = d + 1;
                        fill(distance, maxDist, ll + l, cc + c, direction);
                    }
                    l = 1;
                    if (ll < height - 1 && distance[(ll + l) * width + (cc + c)] > d + 1 && d + 1 <= maxDist) {
                        distance[(ll + l) * width + (cc + c)] = d + 1;
                        fill(distance, maxDist, ll + l, cc + c, direction);
                    }
                }
                break;
        }
    }

    public int[] distanceMatrix(
            final boolean[] bin,
            final boolean distanceTo,
            final int maxDist) {

        int[] distance = getTmpWidthHeightIntArray();

        Arrays.fill(distance, Integer.MAX_VALUE / 2);

        for (int l = 0; l < height; l++) {
            for (int c = 0; c < width; c++) {
                if (bin[l * width + c] == distanceTo) {
                    distance[l * width + c] = 0;
                }
            }
        }

        for (int l = 0; l < height; l++) {
            for (int c = 0; c < width; c++) {
                if (bin[l * width + c] == distanceTo) {
                    fill(distance, maxDist, l, c, 0);
                    fill(distance, maxDist, l, c, 1);
                    fill(distance, maxDist, l, c, 2);
                    fill(distance, maxDist, l, c, 3);
                }
            }
        }

        return distance;
    }

    /**
     * @param img RGB image
     * @param tlx Top-left x pixel, inclusive
     * @param tly Top-left y pixel, inclusive
     * @param brx Bottom-right x pixel, exclusive
     * @param bry Bottom-right y pixel, exclusive
     */
    int[] regionHistogramRGB(int[] img, int tlx, int tly, int brx, int bry, int nBins) {
        int[] hist = new int[(int) pow(nBins, nComp)];

        for (int l = tly; l < bry; l++) {
            for (int c = tlx; c < brx; c++) {
                final int offset = l * width * nComp + c * nComp;
                if (offset < 0) {
                    System.out.println("Error");
                }
                double r = img[offset + 0], g = img[offset + 1], b = img[offset + 2];
                double sum = r + g + b + 1;
                if (sum > 0) {
                    // According to wikipedia the values should be normalized
                    // http://en.wikipedia.org/wiki/Color_histogram
                    r /= sum;
                    g /= sum;
                    b /= sum;
                }
                int rIdx = (int) (r * nBins),
                        gIdx = (int) (g * nBins),
                        bIdx = (int) (b * nBins);
                hist[bIdx * nBins * nBins + gIdx * nBins + rIdx]++;
            }
        }

        return hist;
    }

    public int[] extractHistogramTree(int[] img, Blob blob, int maxDepth, int nBins) {
        return extractHistogramTree(img, blob.getTopLeftX(), blob.getTopLeftY(), blob.getBottomRightX(), blob.getBottomRightY(), maxDepth, 1, nBins);
    }

    public int[] extractHistogramTree(
            int[] img,
            int tlx,
            int tly,
            int brx,
            int bry,
            int maxDepth,
            int depth,
            int nBins) {
        if (depth == maxDepth) {
            return regionHistogramRGB(img, tlx, tly, brx, bry, nBins);
        } else {

            final int size = (int) Math.pow(nBins, 3);

            // 10
            // 00
            int[] hist0 = extractHistogramTree(img, tlx, tly, tlx + (brx - tlx) / 2, tly + (bry - tly) / 2, maxDepth, depth + 1, nBins);
            // 01
            // 00
            int[] hist1 = extractHistogramTree(img, tlx + (brx - tlx) / 2, tly, brx, tly + (bry - tly) / 2, maxDepth, depth + 1, nBins);
            // 00
            // 10
            int[] hist2 = extractHistogramTree(img, tlx, tly + (bry - tly) / 2, tlx + (brx - tlx) / 2, bry, maxDepth, depth + 1, nBins);
            // 00
            // 01
            int[] hist3 = extractHistogramTree(img, tlx + (brx - tlx) / 2, tly + (bry - tly) / 2, brx, bry, maxDepth, depth + 1, nBins);
            int[] hist = new int[size + hist0.length * 4];
            System.arraycopy(hist0, 0, hist, size + 0 * hist0.length, hist0.length);
            System.arraycopy(hist1, 0, hist, size + 1 * hist0.length, hist0.length);
            System.arraycopy(hist2, 0, hist, size + 2 * hist0.length, hist0.length);
            System.arraycopy(hist3, 0, hist, size + 3 * hist0.length, hist0.length);

            for (int i = 0; i < size; i++) {
                hist[i] =
                        (hist[size + hist0.length * 0 + i]
                        + hist[size + hist0.length * 1 + i]
                        + hist[size + hist0.length * 2 + i]
                        + hist[size + hist0.length * 3 + i]) / 4;
            }

            return hist;
        }
    }
    // ==== End Algorithmic Methods ====
}
